#include <drawstuff/drawstuff.h>
#include <ode/ode.h>
#include <fstream>
#include <cmath>
#include <cstdlib>
#include <iostream>

#include "SDLittleDog.h"
#include "SDController.h"
#include "SDUtils.h"
#include "SDKinematics.h"

using namespace std;


// Global pointer to little dog for running drawstuff UI
SDLittleDog *gSimUIDog;
bool gSimRunning;

// Routine for displaying the contents of an SDDogState
void SDDogState::display()
{
	cout << "Dog State:" << endl;
	cout << "pos = [ " 
		 << pos.n[0] << ", " 
		 << pos.n[1] << ", "
		 << pos.n[2] << "]" << endl;
	cout << "ori = [ "
		 << ori.n[0] << ", "
		 << ori.n[1] << ", "
		 << ori.n[2] << "]" << endl;
	cout << "angRates = [ "
		 << angRates.n[0] << ", "
		 << angRates.n[1] << ", "
		 << angRates.n[2] << "]" << endl;
	for ( int i = 0; i < NUM_LEGS; i++ )
	{
		cout << "footPos[" << i << "] = [" 
			 << footPos[i].n[0] << ", "
			 << footPos[i].n[1] << ", "
			 << footPos[i].n[2] << "]" << endl;
	}
	for ( int i = 0; i < NUM_LEGS; i++ )
	{
		cout << "jointAngles[" << i << "] = [" 
			 << jointAngles[i].n[0] << ", "
			 << jointAngles[i].n[1] << ", "
			 << jointAngles[i].n[2] << "]" << endl;
	}
	cout << "time: " << timestamp << endl;		 
}

// constructor
SDLittleDog::SDLittleDog()
{
  mSimulated = true;
  mSimUI = true;
  mInitControl = false;
  mController = NULL;
  for (int i = 0; i < NUM_LEGS; i++) {
    mLegModes[i] = SD_LEG_MODE_PD;
    for (int j = 0; j < NUM_JOINTS; j++) {
      mPDGains[i*3+j][0] = 10.0;
      mPDGains[i*3+j][1] = 0.2;
    }
  }
  mIMUBias = bduVec3f(0.0, 0.0, 0.0);
}

// destructor
SDLittleDog::~SDLittleDog()
{

}

// return the current dog's state
void SDLittleDog::getDogState(SDDogState *dogState)
{
  memcpy(dogState, &mDogState, sizeof(SDDogState));
}

// update the current state of the dog
void SDLittleDog::updateDogState()
{
  bduVec4f qOri;
  bduVec3f eOri;
  bduVec3f acc;
  
  if (mSimulated) {
    // get position, orientation, angular rates, and timestamp
    mSimulator.getDogAngles(mDogState.jointAngles);
    mSimulator.getBodyInformation(LittleDog::B_TRUNK, &mDogState.pos, &qOri);
    SDQuaternionToEuler(qOri, &mDogState.ori);
    mSimulator.getDogAngularRates(&mDogState.angRates);
    mDogState.timestamp = mSimulator.getTimestamp();
  }

  // Either way, use FK now to compute foot positions
  for ( int leg = 0; leg < NUM_LEGS; leg++ )
  {
    bduVec3f localFootPos;
    SDFwdKin(leg, mDogState.jointAngles[leg], &localFootPos);
    SDLocalToWorld(mDogState.pos, mDogState.ori,
		   localFootPos, &mDogState.footPos[leg]);
  }
}


// initialize the simulator
void SDLittleDog::initializeSimulator()
{
  bduVec3f dogPos;
  bduVec4f dogOri;
  bduVec3f dogAngles[NUM_LEGS];
  bduVec3f dogFeet[NUM_LEGS];
  bduVec3f junk;
  int terId;
  bduVec3f terPos;
  bduVec4f terOri;
  int poseType;
  double servoGainMean, servoGainSig, frictionMean, frictionSig;
  int randomSeed;

  // initialize the little dog position from the simulation.params file
  fstream fin("params/simulation.params");
  if (!fin.good()) return;
  
  // read the dog position, orientation, and joint angles or feet position
  fin >> dogPos.n[0] >> dogPos.n[1] >> dogPos.n[2];
  fin >> dogOri.n[0] >> dogOri.n[1] >> dogOri.n[2] >> dogOri.n[3];
  fin >> poseType;
  if (poseType == 1) {
    for (int i = 0; i < NUM_LEGS; i++) {
      for (int j = 0; j < 3; j++) {
	fin >> dogFeet[i].n[j];
      }
      if (i == LittleDog::HL || i == LittleDog::HR) {
	SDInvKin(i, dogFeet[i], &dogAngles[i], &junk);
      } else {
	SDInvKin(i, dogFeet[i], &junk, &dogAngles[i]);
      }
    }
  } else {
    for (int i = 0; i < NUM_LEGS; i++) {
      for (int j = 0; j < NUM_JOINTS; j++) {
	fin >> dogAngles[i].n[j];
      }
    }
  }
    
  mSimulator.setDogPose(dogPos, dogOri, dogAngles);
    
  // set up the terrain boards
  fin >> terId;
  fin >> terPos.n[0] >> terPos.n[1] >> terPos.n[2];
  fin >> terOri.n[0] >> terOri.n[1] >> terOri.n[2] >> terOri.n[3];
  if (terId) {
    mSimulator.setTerrainPose(terPos, terOri);
  }

  // set the servo gain and friction
  fin >> servoGainMean >> servoGainSig >> frictionMean >> frictionSig;
  mSimulator.setServoGain(servoGainMean, servoGainSig);
  mSimulator.setFriction(frictionMean, frictionSig);

  fin >> randomSeed;
  SDSeedRandom(randomSeed);

  // simulate 100 steps to get things going
  mSimulator.setApplyControl(false);
  for (int i = 0; i < 100; i++) {
    mSimulator.step();
  }
  mSimulator.setApplyControl(true);
  
}  


  // start running the trial
LD_ERROR_CODE SDLittleDog::runTrial(SDSimulatorState *state)
{
  mInitialState = state;
  
  if (mSimulated) {
    if (mSimUI) {
      // we can only run one UI at a time
      dsFunctions fn;
      int argc = 2;
      char* argv[2] = {"littledog", "-nopause"};

      fn.version = DS_VERSION;
      fn.start = &dsSimStart;
      fn.step = &dsSimLoop;
      fn.stop = &dsSimStop;
      fn.command = &dsSimCommand;
      fn.path_to_textures = "./textures";
      
      gSimUIDog = this;
      gSimRunning = true;

      dsSimulationLoop(argc, argv, 640, 480, &fn);

    } else {
      // we should add better support for stopping
      if (simStart()) {
        gSimRunning = true;
        while (gSimRunning) {
          simStep();
        }
        simStop();
      }
    }
  }
  return LD_OKAY;
}


// stop the trial
LD_ERROR_CODE SDLittleDog::stopRobot(void)
{
  gSimRunning = false;
  return LD_OKAY;
}


// perform all functions necessary for starting a trial in simulation
bool SDLittleDog::simStart()
{
  return initControl();
}

// take a step in simulation
void SDLittleDog::simStep()
{
  updateControl();
  mSimulator.step();
}

// stop the simulation
void SDLittleDog::simStop()
{
  uninitControl();
}

// draw the simulation
void SDLittleDog::simDraw()
{
  mSimulator.dsDraw();
  mController->dsDraw();
}


// determine whether a given leg hit the ground
bool SDLittleDog::legHit(int leg)
{
  return mSimulator.legHit(leg);
}


// set up the trial
bool SDLittleDog::initControl()
{
    
  mSimulator.setState(mInitialState);
  
  mController->setLittleDog(this);
  mController->reset();
  
  mInitControl = true;
  return true;
}

// update the control for one step
void SDLittleDog::updateControl()
{
  bduVec3f dogAngles[NUM_LEGS];

  // initialize the controller if necessary, or advance the joints
  if (mInitControl) {
    mInitControl = false;
  } else {
    memcpy(&mOldFootForces, &mFootForces, sizeof(bduVec4f));
    mController->advanceAngles();
  }

  // set the desired angles
  updateDogState();
  mController->getAngles(dogAngles);
  if (mSimulated) {
    mSimulator.setDogDesiredAngles(dogAngles);
  }
}

// uninitialize the control
void SDLittleDog::uninitControl()
{
  
}

 

// drawstuff callback at start of simulation - load saved camera
void dsSimStart()
{
  FILE *f = fopen("camera.save", "r");
  static float xyz[3];
  static float hpr[3];
  if (f!=NULL) {
    fread(&xyz, sizeof(float), 3, f);
    fread(&hpr, sizeof(float), 3, f);
    dsSetViewpoint (xyz,hpr);
    fclose(f);
  }

  gSimUIDog->simStart();
}


// drawstuff callback to end simulation
void dsSimStop()
{
  FILE *f = fopen("camera.save", "w");
  float xyz[3], hpr[3];
  dsGetViewpoint(xyz, hpr);
  fwrite(&xyz, sizeof(float), 3, f);
  fwrite(&hpr, sizeof(float), 3, f);
  fclose(f);

  gSimUIDog->simStop();
}

// drawstuff simulation loop callback
void dsSimLoop(int pause)
{
  if (gSimRunning) {
    if (!pause) {
      gSimUIDog->simStep();
    }
  } else {
    dsStop();
  }
  gSimUIDog->simDraw();
}

// drawstuff command callback
void dsSimCommand(int cmd)
{
  bool dynamics;
  switch (cmd) {
  case 'd':
    dynamics = gSimUIDog->getSimulator()->getDynamicsEnabled();
    gSimUIDog->getSimulator()->setDynamicsEnabled(!dynamics);
    break;

  case 't':
    SDSetDrawTerrain(!SDGetDrawTerrain());
    break;

  }
}
